package io.github.iogogogo.commons.db;

import com.google.common.collect.Lists;
import com.google.common.collect.Maps;
import io.github.iogogogo.commons.util.IoUtils;
import io.github.iogogogo.commons.util.SuperStringUtils;
import io.vavr.Tuple;
import io.vavr.Tuple2;
import lombok.Getter;
import lombok.extern.slf4j.Slf4j;

import java.sql.*;
import java.util.Arrays;
import java.util.List;
import java.util.Map;
import java.util.function.Consumer;
import java.util.stream.Collectors;


/**
 * Created by tao.zeng on 2024/11/7.
 */
@Slf4j
public class DBHelper {

    @Getter
    public enum DefaultDriverEnum {

        /**
         * sqlite
         */
        _DRIVER_SQLITE("sqlite", "org.sqlite.JDBC"),

        /**
         * MySQL
         */
        _DRIVER_MYSQL("MySQL", "com.mysql.cj.jdbc.Driver");

        private final String database;

        private final String driver;

        DefaultDriverEnum(String database, String driver) {
            this.database = database;
            this.driver = driver;
        }
    }

    /**
     * jdbc连接对象
     */
    private final Connection conn;

    /**
     * SQL处理对象
     */
    private PreparedStatement ps;


    public DBHelper(String driver, String url) throws ClassNotFoundException, SQLException {
        this(driver, url, null, null);
    }

    /**
     * @param driver   数据库驱动
     * @param url      数据库连接字符串
     * @param username 用户名
     * @param password 密码
     * @throws ClassNotFoundException ClassNotFoundException
     * @throws SQLException           SQLException
     */
    public DBHelper(String driver, String url, String username, String password) throws ClassNotFoundException, SQLException {
        Class.forName(driver);
        conn = DriverManager.getConnection(url, username, password);
    }

    /**
     * SQL参数注入
     *
     * @param sql   sql
     * @param param 参数
     * @throws SQLException SQL参数注入异常
     */
    private void inject(String sql, List<Object> param) throws SQLException {
        if (ps != null) ps.close();
        ps = conn.prepareStatement(sql);
        if (log.isDebugEnabled()) log.debug("inject sql ===> {}", sql);
        for (int i = 0; i < param.size(); i++) {
            Object obj = param.get(i);
            ps.setObject(i + 1, obj);
            log.debug("param[{}] ===> {}", i, obj);
        }
    }

    /**
     * 批量插入数据
     *
     * @param sql      SQL
     * @param data     需要插入的数据
     * @param consumer 每一条consumer处理对象
     * @param <T>      插入的数据类型
     * @return 每一批插入的结果
     */
    public <T> List<int[]> batchSave(String sql, List<T> data, Consumer<Tuple2<T, PreparedStatement>> consumer) {
        return batchSave(sql, 500, data, consumer);
    }

    /**
     * 批量插入数据
     * <blockquote>
     * <pre>
     * String url = String.join("", "jdbc:sqlite:", "sqlite.db");
     *
     * DBHelper dbHelper = new DBHelper(DBHelper.DefaultDriverEnum._DRIVER_SQLITE.getDriver(), url);
     * List<User> users = new ArrayList<>(350000);
     *
     * String sql = "INSERT INTO t_user ( username, age) VALUES ( ?, ?)";
     * List<int[]> resultList = dbHelper.batchSave(sql, 1000, users, x -> {
     *     try {
     *         User user = x._1;
     *         PreparedStatement ps = x._2;
     *         ps.setObject(1, user.getUsername());
     *         ps.setObject(2, user.getAge());
     *     } catch (SQLException e) {
     *         throw new RuntimeException(e);
     *     }
     * });
     * log.info("resultList:{}", resultList.size());
     * </pre>
     * </blockquote>
     *
     * @param sql       SQL
     * @param batchSize 每一批大小
     * @param data      需要插入的数据
     * @param consumer  每一条consumer处理对象
     * @param <T>       插入的数据类型
     * @return 每一批插入的结果
     */
    public <T> List<int[]> batchSave(String sql, int batchSize, List<T> data, Consumer<Tuple2<T, PreparedStatement>> consumer) {
        List<int[]> result = Lists.newArrayList();
        try {
            if (log.isDebugEnabled()) log.debug("batchSave sql ===> {}", sql);
            ps = conn.prepareStatement(sql);

            // 关闭自动提交事务，改为手动提交
            conn.setAutoCommit(false);
            // 数据分片
            List<List<T>> lists = Lists.partition(data, batchSize);

            // 循环处理每一个分片的数据
            for (List<T> list : lists) {
                for (T t : list) {
                    // 绑定参数
                    consumer.accept(Tuple.of(t, ps));
                    // 添加到批处理中
                    ps.addBatch();
                    log.debug("param ===> {}", t);
                }
                // 执行一批数据插入
                int[] batch = ps.executeBatch();
                // 手动提交事务
                conn.commit();

                result.add(batch);
            }
            return result;
        } catch (SQLException e) {
            throw new RuntimeException(e);
        } finally {
            this.close();
        }
    }

    /**
     * 执行更新
     *
     * @param sql   sql
     * @param param 参数
     * @return 更新行数
     */
    public int update(String sql, Object... param) {
        return update(sql, Arrays.stream(param).collect(Collectors.toList()));
    }

    /**
     * 执行更新
     *
     * @param sql   sql
     * @param param 参数
     * @return 更新行数
     */
    public int update(String sql, List<Object> param) {
        try {
            inject(sql, param);
            return ps.executeUpdate();
        } catch (SQLException e) {
            throw new RuntimeException(e);
        } finally {
            this.close();
        }
    }

    /**
     * 执行查询
     *
     * @param sql   sql
     * @param param 参数
     * @return 查询结果
     */
    public ResultSet query(String sql, Object... param) throws SQLException {
        return query(sql, Arrays.stream(param).collect(Collectors.toList()));
    }

    /**
     * 执行查询
     *
     * <blockquote>
     * <pre>
     * String sql = "....";
     * try (ResultSet rs = dbHelper.query(sql)) {
     *     List<Map<String, Object>> mapList = dbHelper.rsToListMap(rs, false);
     *     log.info("mapList:{}", mapList.size());
     * }
     * </pre>
     * </blockquote>
     *
     * @param sql   sql
     * @param param 参数
     * @return 查询结果
     */
    public ResultSet query(String sql, List<Object> param) throws SQLException {
        inject(sql, param);
        return ps.executeQuery();
    }

    /**
     * rs转换为ListMap对象
     *
     * @param rs ResultSet
     * @return 转换结果集
     * @throws SQLException SQLException
     */
    public List<Map<String, Object>> rsToListMap(ResultSet rs) throws SQLException {
        return rsToListMap(rs, true);
    }

    /**
     * rs转换为ListMap对象
     *
     * @param rs               ResultSet
     * @param enableLineToHump 是否启用下划线转map
     * @return 转换结果集
     * @throws SQLException SQLException
     */
    public List<Map<String, Object>> rsToListMap(ResultSet rs, boolean enableLineToHump) throws SQLException {
        ResultSetMetaData metaData = rs.getMetaData();
        int columnCount = metaData.getColumnCount();

        List<Map<String, Object>> mapList = Lists.newArrayList();
        while (rs.next()) {
            Map<String, Object> map = Maps.newHashMap();
            for (int i = 1; i <= columnCount; i++) {
                String key = metaData.getColumnName(i);
                String val = rs.getString(i);
                // 下划线转驼峰
                map.put(enableLineToHump ? SuperStringUtils.lineToHump(key) : key, val);
            }
            mapList.add(map);
        }

        return mapList;
    }

    /**
     * 释放数据库连接资源
     *
     * @param rs 关闭 ResultSet
     */
    public void close(ResultSet rs) throws SQLException {
        IoUtils.close(rs, ps, conn);
    }

    /**
     * 释放数据库连接资源
     */
    public void close() {
        try {
            this.close(null);
        } catch (SQLException e) {
            throw new RuntimeException(e);
        }
    }
}
